open Base
open Math

type dir = [`H | `V]
type fig = Rect of Dim.t * Color.t

type tree = (arm, fig) Generic_tree.t
and arm = [dir | `Capture of Event.Id.t]
and t = tree

let arm dir ts : t =
  Arm((dir :> arm), ts) [@@ocaml.inline always]
let figure fg : t =
  Leaf(fg) [@@ocaml.inline always]
let capture cid t : t =
  Arm(`Capture(cid), [t]) [@@ocaml.inline always]

module Zipper = Generic_tree.Zipper

include Generic_tree.Make_show_eq
    (struct
      type t = arm
      let to_string = function
        | `H -> "H"
        | `V -> "V"
        | `Capture(eid) -> Printf.sprintf "cap{%s}"
                             (Event.Id.to_string eid)
      let equal =
        Polymorphic_compare.(=)
    end)
    (struct
      type t = fig
      let to_string = function
        | Rect(dim, color) ->
           Printf.sprintf "Rect(%s,%s)"
             (Dim.to_string dim)
             (Color.to_string color)
      let equal a b = match a, b with
        | Rect(d1, c1), Rect(d2, c2) ->
           Dim.equal d1 d2 && Color.equal c1 c2
    end)


(*** paths ***)

module Path = struct
  type t = int list

  let root = []

  let of_list_rev = List.rev

  let to_string p =
    Printf.sprintf "path[%s]" @@
    String.concat ~sep:";" @@
    List.map p ~f:Int.to_string

  include Comparable.Make(struct
      type nonrec t = t

      let sexp_of_t p =
        Sexp.List(List.map p ~f:(fun x -> Sexp.Atom(Int.to_string x)))

      let rec compare p1 p2 = match p1, p2 with
        | _,      []     -> - List.length p1
        | [],     _      -> List.length p2
        | i::p1', j::p2' ->
           if i = j then
             compare p1' p2'
           else
             Int.compare i j
    end)
end

(*** updating ***)

type update = Path.t * fig

let apply_updates_exn ups t : t =
  Zipper.to_tree @@ snd @@
  Sequence.fold ups ~init:(Path.root, Zipper.of_tree t)
    ~f:(fun (p, zip) (p', fig) ->
        let zip = Zipper.set (Leaf(fig))
                    (Zipper.nav_exn ~prev:p ~targ:p' zip) in
        (p', zip))


(*** generic geometry calculation ***)

module type GEOMETRY_MONOID = sig
  type t
  val concat : t Sequence.t -> t
  val rect : AABB.t -> Color.t -> t
  val capture : Event.Id.t -> t -> t
end

let geometry
      (type a)
      (module A : GEOMETRY_MONOID with type t = a)
      (tree : t) : a =

  let rec compute : t -> (int * int * (AABB.t -> a)) = function
    | Leaf(Rect({ w ; h } as dim, fill)) ->
       (w, h, fun bounds -> A.rect (AABB.center bounds dim) fill)

    | Arm(`Capture(cid), ts) ->
       let (w, h, go) = compute (List.hd_exn ts) in
       (w, h, fun bounds -> A.capture cid (go bounds))

    | Arm((`H | `V) as dir, ts) ->
       let axis1 x y = match dir with `H -> x | `V -> y in
       let axis2 x y = match dir with `H -> y | `V -> x in
       let pair a1 a2 = match dir with `H -> (a1, a2) | `V -> (a2, a1) in

       let (sizes_funcs, max_size, total_size) =
         List.fold_right ts ~init:([], 0, 0)
           ~f:(fun t (s_f, max_sz, tot_sz) ->
               let (w, h, func) = compute t in
               ((axis1 w h, func)::s_f,
                max (axis2 w h) max_sz,
                axis1 w h + tot_sz))
       in

       let go {bx;by;bh;bw} =
         let init = (axis1 bw bh - total_size) / 2 in
         Sequence.folding_map (Sequence.of_list sizes_funcs) ~init
           ~f:(fun off (size, func) ->
               (off + size,
                func { bx = bx + axis1 off 0
                     ; by = by + axis2 off 0
                     ; bw = axis1 size bw
                     ; bh = axis2 size bh }))
         |> A.concat
       in

       let (w, h) = pair total_size max_size in
       (w, h, go)
  in
  let (bw, bh, func) = compute tree in
  func {bx=0;by=0;bw;bh}


(*** rendering ***)

module Draw = struct
  type t = Rect of AABB.t * Color.t

  let equal x y = match x, y with
    | Rect(b1, c1), Rect(b2, c2) ->
       AABB.equal b1 b2 && Color.equal c1 c2

  let to_string = function
    | Rect(b, c) -> Printf.sprintf "Rect(%s, %s)"
                      (AABB.to_string b)
                      (Color.to_string c)
end

module Draw_monoid = struct
  type nonrec t = Draw.t Sequence.t
  let concat = Sequence.concat
  let rect aabb color = Sequence.singleton (Draw.Rect(aabb, color))
  let capture _ s = s
end

let render tree : Draw.t Sequence.t =
  geometry (module Draw_monoid) tree


(*** collision ***)

module Collide_monoid = struct
  type t = Pos.t -> Event.Id.t list option
  let concat ts pos = Sequence.find_map ts ~f:(fun f -> f pos)
  let rect bounds _ pos = Option.some_if (AABB.contains bounds pos) []
  let capture cid f pos = Option.map (f pos) ~f:(List.cons cid)
end

let capture_point ~pos tree : Event.Id.t list =
  Option.value (geometry (module Collide_monoid) tree pos)
    ~default:[]
